#include "sim.hpp"
#include "../rescue/shuf.hpp"
#include <algorithm>

using namespace omelette;

Oracle::Oracle(const State &s) : ops(omelette::solve(s)) {
}

Op Oracle::plan(const State&) {
	if (ops.empty())
		return Op::Nop;

	auto op = ops.back();
	ops.pop_back();
	return op;
}

Hindsight::Hindsight(unsigned int n) : nsamp(n) {
}

Op Hindsight::plan(const State &state) {
	if (state.isgoal())
		return Op::Nop;

	std::vector<State> samps;
	for (unsigned int i = 0; i < nsamp; i++)
		samps.push_back(state.future());

	double mincost = std::numeric_limits<double>::infinity();
	Op minop = Op::Nop;

	auto ops = state.ops();
	shuf(ops);

	for (auto op : ops) {
		fprintf(stderr, "\tPlanning for %s", op.str().c_str());

		double t = walltime();
		double e = 0;
		for (auto s : samps) {
			auto path = omelette::solve(s.act(op));
			e += path.size();
			fprintf(stderr, ".");
		}
		e /= samps.size();
		fprintf(stderr, "%g (%gs)\n", e, walltime()-t);

		if (e < mincost) {
			mincost = e;
			minop = op;
		}
	}

	fprintf(stderr, "\tReturning %s (%g)\n", minop.str().c_str(), mincost);
	return minop;
}

Op Handmade::plan(const omelette::State &s) {
	Op op;

	assert (stillzero || !s.bowls[0].bad);

	if (s.bowls[0].n == 3 && s.bowls[0].known) {
		assert (!s.bowls[0].bad);
		return Op::Nop;
	}

	switch (state) {
	case Grab: {
		assert (s.hand == Empty);

		// If we've dirtied all of the bowls with bad eggs
		// then we must clean one to have somewhere
		// to break.
		int i ;
		for (i = NBowls-1; i > 0; i--) {
			if (s.bowls[i].n == 0)
				break;
		}
		if (i <= 0) {
			assert (s.bowls[NBowls-1].known && s.bowls[NBowls-1].bad);	
			op = Op::clean(NBowls-1);
			fprintf(stderr, "state=Grab, no empty bowls, cleaning %d\n",
				NBowls-1);
			break;
		}

		// There is a clean bowl into which we can break,
		// grab the next egg.
		op = Op::grab();
		state = BreakSmell;
		fprintf(stderr, "state=Grab: grabbing an egg\n");
		break;
	}

	case BreakSmell: {
		if (s.hand != Empty) {	
			int i = 0;
			if (s.bowls[0].n > 0)
				i++;
			for (; i < NBowls; i++) {
				if (s.bowls[i].n == 0)
					break;
			}
			assert (i < NBowls);
			op = Op::breakegg(i);
			fprintf(stderr, "state=Break: breaking into %d\n", i);
			state = BreakSmell;
			break;
		}

		// We just broke an egg into the greatest index
		// bowl with something in it.  We need to smell
		// the bowl to see if it was any good.
		int i;
		for (i = NBowls-1; i >= 0; i--) {
			if (s.bowls[i].n > 0)
				break;
		}
		assert(i >= 0);
		op = Op::smell(i);
		fprintf(stderr, "state=Smell: smelling %d\n", i);
		state = Pour;
		break;
	}

	case Pour: {
		// We just smelt a bowl, if it is not bad then we should
		// pour it into bowl 0, otherwise we should grab the
		// next egg.
		int i;
		for (i = NBowls-1; i >= 0; i--) {
			if (s.bowls[i].n > 0)
				break;
		}
		assert(i >= 0);
		assert (s.bowls[i].known);
		state = Grab;
		if (!s.bowls[i].bad && i != 0) {
			stillzero = false;
			op = Op::pour(i, 0);
			fprintf(stderr, "state=Pour: good egg, pouring %d into 0\n", i);
		} else if (s.bowls[i].bad && i == 0) {
			fprintf(stderr, "state=Pour: bad egg in %d: clean\n", i);
			op = Op::clean(0);
		} else {
			fprintf(stderr, "state=Pour: leaving %d: next\n", i);
			op = plan(s);
		}
		break;
	}

	}

	bool found = false;
	for (auto o : s.ops()) {
		if (o == op) {
			found = true;
			break;
		}
	}
	assert (found);
	return op;
}

void sim(const State &init, Planner *planner) {
	State cur(init);
	std::vector<Op> ops;

	double stepmin = std::numeric_limits<double>::infinity();
	double stepmax = -std::numeric_limits<double>::infinity();
	double simstart = walltime();

	dfrowhdr(stdout, "step", 2, "number", "wall time");

	for (;;) {
		fprintf(stderr, "%s\n", cur.str().c_str());

		double stepstart = walltime();

		auto op = planner->plan(cur);

		fprintf(stderr, "%s\n", op.str().c_str());

		if (op.kind == Op::None) {
			assert (cur.isgoal());
			break;
		}

		double steptime = walltime() - stepstart;
		stepmin = std::min(steptime, stepmin);
		stepmax = std::max(steptime, stepmax);

		cur = cur.act(op);
		ops.push_back(op);

		dfrow(stdout, "step", "ug", (unsigned long) ops.size(), steptime);
	}

	dfpair(stdout, "total wall time", "%g", walltime() - simstart);
	dfpair(stdout, "final sol length", "%lu", (unsigned long) ops.size());
	dfpair(stdout, "max step time", "%g", stepmax);
	dfpair(stdout, "min step time", "%g", stepmin);
}